

package com.hmc.common.interceptor;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.Map;


/**
 * 数据过滤
 *
 * 数据过滤拦截器，用于在查询前根据数据权限范围修改SQL。
 * @author
 */
public class DataFilterInterceptor implements InnerInterceptor {

    /**
     * 在查询执行前，根据数据权限范围修改SQL。
     *
     * @param executor 执行器
     * @param ms 映射语句
     * @param parameter 参数对象
     * @param rowBounds 行限制
     * @param resultHandler 结果处理器
     * @param boundSql 绑定的SQL对象
     */
    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
        // 获取数据权限范围
        DataScope scope = getDataScope(parameter);
        // 如果没有数据权限范围或范围条件为空，则不进行数据过滤
        if(scope == null || StrUtil.isBlank(scope.getSqlFilter())){
            return;
        }

        // 根据数据权限范围拼接新的SQL
        String buildSql = getSelect(boundSql.getSql(), scope);

        // 重写SQL
        PluginUtils.mpBoundSql(boundSql).sql(buildSql);
    }

    /**
     * 从参数中获取数据权限范围对象。
     *
     * @param parameter 参数对象
     * @return 数据权限范围对象，如果不存在则返回null
     */
    private DataScope getDataScope(Object parameter){
        if (parameter == null){
            return null;
        }

        // 尝试从参数中寻找DataScope对象
        if (parameter instanceof Map) {
            Map<?, ?> parameterMap = (Map<?, ?>) parameter;
            for (Map.Entry entry : parameterMap.entrySet()) {
                if (entry.getValue() != null && entry.getValue() instanceof DataScope) {
                    return (DataScope) entry.getValue();
                }
            }
        } else if (parameter instanceof DataScope) {
            return (DataScope) parameter;
        }

        return null;
    }

    /**
     * 根据数据权限范围修改SQL。
     *
     * @param buildSql 原始SQL
     * @param scope 数据权限范围
     * @return 修改后的SQL
     */
    private String getSelect(String buildSql, DataScope scope){
        try {
            // 解析SQL并修改条件以添加数据权限范围限制
            Select select = (Select) CCJSqlParserUtil.parse(buildSql);
            PlainSelect plainSelect = (PlainSelect) select.getSelectBody();

            Expression expression = plainSelect.getWhere();
            // 如果原始SQL无WHERE条件，则直接添加数据权限范围条件
            if(expression == null){
                plainSelect.setWhere(new StringValue(scope.getSqlFilter()));
            }else{
                // 如果原始SQL已有WHERE条件，则将数据权限范围条件与原有条件用AND连接
                AndExpression andExpression =  new AndExpression(expression, new StringValue(scope.getSqlFilter()));
                plainSelect.setWhere(andExpression);
            }

            // 生成并返回修改后的SQL
            return select.toString().replaceAll("'", "");
        }catch (JSQLParserException e){
            // 如果解析SQL失败，则返回原始SQL
            return buildSql;
        }
    }
}
